can we have a to(device) API? before this call, all the prepared input lives in cpu.

So you suggest like the following?

        builder = ModelInputForGPUBuilder(...)
        for seq_group_metadata in seq_group_metadata_list:
            builder.add_seq_group(seq_group_metadata)
        model_input = builder.build(self) # All tensors are on CPU.
        return model_input.to(device) # Move to GPU.

This would introduce some overheads as we create tensors on CPU and then move to GPU, instead of directly creating them on GPU.

yes, this would simplify the control plane communication.

actually, if you directly creating them on GPU, pytorch still needs to prepare them in cpu, and then move it to GPU together. I don't think there would be significant overhead here.

The following code snippets are essentially the same:

data = list(range(1000))
gpu_tensor = torch.tensor(data, dtype=torch.int, device="cuda")

data = list(range(1000))
cpu_tensor = torch.tensor(data, dtype=torch.int, device="cpu")
gpu_tensor = cpu_tensor.to("cuda")

In the first code, pytorch still needs to create a cpu array to hold the Python data structure, and then move it to GPU.

If we do the cpu-gpu movement before we broadcast it to the rest workers, they will wait for the broadcast.

timeline:

driver worker: |  prepare input in cpu | move to GPU | broadcast |
rest workers:  |    wait for broadcast                         |

If we do the cpu-gpu movement after we broadcast, the driver worker sends cpu data (which can be fast because we have shared memory transport), and all workers move data to gpu respectively.

timeline:

driver worker: |  prepare input in cpu | broadcast | move to GPU |
rest workers:  |    wait for broadcast           | move to GPU |

Because these data are usually quite small, I suppose the second approach would be faster. Broadcasting small cpu data would be easier. Broadcasting small GPU data with NCCL is not good in performance.

Hmm...make sense since we will have a Python list anyways. I'll think about that and add this API if no other concerns. Thanks.

I've basically done the first round of refactoring and feel this is doable, but I'd prefer to defer this until #6052 is merged, because I plan to add .to() to the following data classes:

1. ModelInputForGPU, which will implicitly call attn_metadata.to(...) (the next point)
2. AttentionMetadata <- conflict

why is compute_slot_mapping_start_idx splitted from compute_slot_mapping ? Seems like they are used together?

Yeah the only reason is if we combine these 2 then the argument list would be more ugly:

compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,
    seq_len, context_len,
    self.block_size,
    seq_group_metadata.block_tables,
    is_prompt, query_len, context_len, self.sliding_window, self.use_v2_block_manager)

I don't have preference tho so if you prefer to combine them I could do that. Plz let me know

I prefer to combine it if you don't mind! Agreed there are too many args, but I think it is better than having 2 calls depending on each other (imo it is harder to use). But it is nit!

Sure. I'll merge this one and do it in the next PR.